#!/usr/bin/env python
# -*- coding: utf -*-
# from MLab import *
# from FFT import *
from dislin import *
from numpy import *
import struct
import sys
import wave
from functions import everyOther

# Input args
lin_log = sys.argv[3]
window_name = sys.argv[4]   


# open the wave file
fp = wave.open(sys.argv[1], "rb")

(nchannels, sampwidth, sample_rate, total_num_samps, comptype, compname) = fp.getparams ()

detailed = 0
window_size = int(sys.argv[2])
if detailed == 0:
    num_fft = (total_num_samps / window_size) - 2
else:
    num_fft = total_num_samps - window_size

print "Sampling frequency: " + str(sample_rate)
print "Total number of samples: " + str(total_num_samps)
print "FFT size: " + str(window_size)
print "Number of FFTs performed: " + str(num_fft)
print "Used window: " + str(window_name)
if lin_log == 'log':
    print "Using logarythmic scale: yes"
else:
    print "Using logarythmic scale: no"

# create temporary working array
temp = zeros((num_fft, window_size), float)
if detailed == 1: 
    signal = fp.readframes(total_num_samps)
    if nchannels == 2:
        signal = array (list (everyOther (signal, 1)))

# read in the data from the file
for i in range(num_fft):
    if detailed == 1:
        tempb = signal[i:i + window_size]
    else:        
        tempb = fp.readframes(window_size)  # signal[i:i + window_size]  # fp.readframes(window_size);
        if nchannels == 2:
            tempb = array (list (everyOther (tempb, 1)))
    temp[i, :] = array(struct.unpack("%dB" % (window_size), tempb[0:window_size]), float) - 128.0
fp.close()

# Window the data plus other possible windows
if window_name == 'hamming':
        temp = temp * hamming(window_size)
else:
    if window_name == 'hanning':
        temp = temp * hanning(window_size)
    else: 
        if window_name == 'kaiser':
            temp = temp * kaiser(window_size, 0.5)

# Transform with the FFT, Return Power
freq_pwr = 2 * abs(fft.rfft(temp, window_size)) / window_size
for i in range(num_fft):
    freq_pwr[i, 0] *= 2
if lin_log == 'log':
    max_val = max([max(l) for l in freq_pwr])
    freq_pwr = 10 * log10(1e-20 + freq_pwr / max_val)

# Plot the result
gradient_res = 512
max_val = max([max(l) for l in freq_pwr])
min_val = min([min(l) for l in freq_pwr])
print "Min val: " + str(min_val)
print "Max val: " + str(max_val)
n_out_pts = (window_size / 2) + 1
y_axis = 0.5 * float(sample_rate) / n_out_pts * arange(n_out_pts)
x_axis = (total_num_samps / float(sample_rate)) / num_fft * arange(num_fft)
gradiens = arange(min_val, max_val, (max_val - min_val) / gradient_res)
setvar('X', "Time (sec)")
setvar('Y', "Frequency (Hertz)")
conshade(freq_pwr, x_axis, y_axis, gradiens)
disfin()
